#include "ps/ps.h"

namespace ps {
namespace tensorflow {
/**
 * @brief 用于tensorflow端的ps::KVWorker生命周期管理RAII类，无功能逻辑
 * \see ps::KVWorker
 *
 * @tparam T
 */
template<class T>
class PSWorker
{
  public:
    PSWorker(int customer_id = 0, int appid = 0)
      : _customer_id(customer_id)
    {
        ps::Start(customer_id);
        _worker = std::make_shared<ps::KVWorker<T>>(customer_id, appid);
    }

    virtual ~PSWorker() { ps::Finalize(_customer_id); }

    std::shared_ptr<ps::KVWorker<T>> worker() { return _worker; }

  private:
    std::shared_ptr<ps::KVWorker<T>> _worker;
    int _customer_id;
};

/**
 * @brief 获取KVWorker实例
 *
 * @return KVWorker实例共享指针(std::shared_ptr)
 * \see ps::KVWorker
 * \see std::shared_ptr
 */
static std::shared_ptr<ps::KVWorker<float>>
worker()
{
    static PSWorker<float> ps(0, 0);
    return ps.worker();
}


/**
 * @brief  将参数push到参数服务器，一般来说只有rank=0节点启动时会将参数push一次
 * @param keys 参数key，计算方式为::tensorflow::hash(var.name)
 * @param data flat后的参数
 * @param lens 参数值flaten后的长度
 */
void
ps_push_data(std::vector<uint64_t> const& keys,
             std::vector<float> const& data,
             std::vector<int> const& lens)
{
    try {
        auto w = worker();
        w->Wait(w->Push(keys, data, lens));
    } catch (dmlc::Error const& e) {
        LOG(ERROR) << "Push data error. what: " << e.what();
    }
}

/**
 * @brief 从参数服务器获取参数,
 * 一般来说只有非rank=0节点启动时拉取一次参数，更新阶段使用
 * ps::tensorflow::ps_push_pull_data 来更新参数
 *
 * @param keys 参数key，计算方式为::tensorflow::hash(var.name)
 * @param data flat后的参数
 * @param lens 参数值flaten后的长度
 */
void
ps_pull_data(std::vector<uint64_t> const& keys,
             std::vector<float>& data,
             std::vector<int>& lens)
{
    auto w = worker();
    try {
        w->Wait(w->Pull(keys, &data, &lens));
    } catch (dmlc::Error const& e) {
        LOG(ERROR) << "Pull data error. what: " << e.what();
    }
}

/**
 * @brief 推送梯度到参数服务器，并pull更新后的参数
 *
 * @param keys 参数key
 * @param data flat后的梯度矩阵
 * @param out  flat后的最新参数
 * @param lens 梯度矩阵flat后的长度
 */
void
ps_push_pull_data(std::vector<uint64_t> const& keys,
                  std::vector<float> const& data,
                  std::vector<float>& out,
                  std::vector<int>& lens)
{
    auto w = worker();
    try {
        // w->Wait(w->Push(keys, data, lens));
        w->Wait(w->PushPull(keys, data, &out, &lens));
    } catch (dmlc::Error const& e) {
        LOG(ERROR) << "Push data error. what: " << e.what();
    }
}

int ps_my_rank()
{
    volatile auto w = worker();
    (void)w;
    return ps::MyRank();
}
}
}

